#未使用pysaprk的最后能跑的版本

from flask import Blueprint, request, jsonify, render_template
from werkzeug.utils import secure_filename
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np
import time
import os
import csv
from datetime import datetime
import pymysql

train_bp = Blueprint('train', __name__)

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'Ab_data',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}


def get_db_connection():
    """创建数据库连接"""
    return pymysql.connect(**DB_CONFIG)


def create_new_abalone_table():
    """创建新表（如果不存在）"""
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            # 检查表是否存在
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS new_abalone (
                    id INT PRIMARY KEY,
                    sex VARCHAR(10),
                    length FLOAT,
                    diameter FLOAT,
                    height FLOAT,
                    whole_weight FLOAT,
                    shucked_weight FLOAT,
                    viscera_weight FLOAT,
                    shell_weight FLOAT,
                    rings INT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                )
            """)
        connection.commit()
    finally:
        connection.close()


# 全局变量存储当前模型
current_model = None
model_type = "LR"
X_train, X_test, y_train, y_test = None, None, None, None

# 确保表存在
create_new_abalone_table()


@train_bp.route('/train', methods=['GET'])
def train_page():
    return render_template('train.html')


# ... (保持之前的导入和配置不变)

@train_bp.route('/api/train/add_data', methods=['POST'])
def add_data():
    try:
        data = request.form
        print(data)
        # 验证必填字段
        required_fields = ['sex', 'length', 'diameter', 'height',
                           'whole_weight', 'shucked_weight',
                           'viscera_weight', 'shell_weight', 'rings']
        for field in required_fields:
            if field not in data or not data[field]:
                return jsonify({'success': False, 'message': f'缺少必填字段: {field}'}), 400

        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                # 获取当前最大ID
                cursor.execute("SELECT MAX(id) AS max_id FROM new_abalone")
                result = cursor.fetchone()
                new_id = (result['max_id'] or 0) + 1

                # 转换性别
                sex_str = 'M' if data['sex'] == '1' else ('F' if data['sex'] == '0' else 'I')

                sql = """
                INSERT INTO new_abalone 
                (id, sex, length, diameter, height, whole_weight, 
                 shucked_weight, viscera_weight, shell_weight, rings)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """
                cursor.execute(sql, (
                    new_id,
                    sex_str,
                    float(data['length']),
                    float(data['diameter']),
                    float(data['height']),
                    float(data['whole_weight']),
                    float(data['shucked_weight']),
                    float(data['viscera_weight']),
                    float(data['shell_weight']),
                    int(data['rings'])
                ))
            connection.commit()
            return jsonify({
                'success': True,
                'message': '数据添加成功',
                'id': new_id
            })
        except pymysql.Error as e:
            connection.rollback()
            return jsonify({
                'success': False,
                'message': f'数据库错误: {str(e)}'
            }), 500
        finally:
            connection.close()
    except Exception as e:
        return jsonify({
            'success': False,
            'message': f'服务器错误: {str(e)}'
        }), 500
@train_bp.route('/api/train/get_single_data/<int:id>', methods=['GET'])
def get_single_data(id):
    try:
        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                sql = """
                SELECT id, sex, length, diameter, height, 
                       whole_weight, shucked_weight, 
                       viscera_weight, shell_weight, rings
                FROM new_abalone
                WHERE id = %s
                """
                cursor.execute(sql, (id,))
                record = cursor.fetchone()

                if not record:
                    return jsonify({'success': False, 'message': '记录不存在'}), 404

                # 转换sex为前端使用的格式
                record['sex'] = '1' if record['sex'] == 'M' else ('0' if record['sex'] == 'F' else '-1')
                return jsonify({
                    'success': True,
                    'record': record
                })
        finally:
            connection.close()
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500


@train_bp.route('/api/train/update_data', methods=['POST'])
def update_data():
    try:
        data = request.form

        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                sql = """
                UPDATE new_abalone 
                SET sex = %s,
                    length = %s,
                    diameter = %s,
                    height = %s,
                    whole_weight = %s,
                    shucked_weight = %s,
                    viscera_weight = %s,
                    shell_weight = %s,
                    rings = %s
                WHERE id = %s
                """
                cursor.execute(sql, (
                    'M' if data['sex'] == '1' else ('F' if data['sex'] == '0' else 'I'),
                    float(data['length']),
                    float(data['diameter']),
                    float(data['height']),
                    float(data['whole_weight']),
                    float(data['shucked_weight']),
                    float(data['viscera_weight']),
                    float(data['shell_weight']),
                    int(data['rings']),
                    int(data['id'])
                ))
            connection.commit()

            return jsonify({'success': True, 'message': '数据更新成功'})
        finally:
            connection.close()
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 400

@train_bp.route('/api/train/batch_import', methods=['POST'])
def batch_import():
    try:
        if 'csvFile' not in request.files:
            return jsonify({'success': False, 'message': '没有上传文件'}), 400

        file = request.files['csvFile']
        if file.filename == '':
            return jsonify({'success': False, 'message': '没有选择文件'}), 400

        if file and file.filename.endswith('.csv'):
            filename = secure_filename(file.filename)
            file_path = os.path.join('uploads', filename)
            os.makedirs('uploads', exist_ok=True)
            file.save(file_path)

            connection = get_db_connection()
            try:
                with connection.cursor() as cursor:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        reader = csv.DictReader(f)
                        count = 0
                        for row in reader:
                            try:
                                # 转换sex字段
                                sex_value = row.get('sex', '-1')
                                sex_mapping = {'1': 'M', '0': 'F', '-1': 'I'}
                                sex_str = sex_mapping.get(sex_value, 'I')  # 默认I

                                # 转换age/rings字段（根据实际CSV字段名）
                                #rings = int(row.get('age') or row.get('rings', 0))

                                sql = """
                                INSERT INTO new_abalone 
                                (id, sex, length, diameter, height, 
                                 whole_weight, shucked_weight, 
                                 viscera_weight, shell_weight, rings)
                                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                                """
                                cursor.execute(sql, (
                                    int(row['id']),
                                    sex_str,
                                    float(row['length']),
                                    float(row['diameter']),
                                    float(row['height']),
                                    float(row['whole_weight']),
                                    float(row['shucked_weight']),
                                    float(row['viscera_weight']),
                                    float(row['shell_weight']),
                                    rings
                                ))
                                count += 1
                            except Exception as e:
                                print(f"Error processing row {row}: {str(e)}")
                                continue

                    connection.commit()
                    os.remove(file_path)
                    return jsonify({
                        'success': True,
                        'message': f'成功导入 {count} 条数据',
                        'count': count
                    })
            finally:
                connection.close()
        else:
            return jsonify({'success': False, 'message': '仅支持CSV文件'}), 400
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500

# ... (保持之前的导入和配置不变)

@train_bp.route('/api/train/get_data', methods=['GET'])
def get_data():
    try:
        page = int(request.args.get('page', 1))
        per_page = int(request.args.get('per_page', 10))  # 添加每页显示数量参数

        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                # 查询数据
                offset = (page - 1) * per_page
                sql = """
                SELECT id, sex, length, diameter, height, 
                       whole_weight, shucked_weight, 
                       viscera_weight, shell_weight, rings
                FROM new_abalone
                ORDER BY id
                LIMIT %s OFFSET %s
                """
                cursor.execute(sql, (per_page, offset))
                records = cursor.fetchall()

                # 查询总数
                cursor.execute("SELECT COUNT(*) AS total FROM new_abalone")
                total_result = cursor.fetchone()
                total = total_result['total'] if total_result else 0

                # 计算总页数
                total_pages = (total + per_page - 1) // per_page if total > 0 else 1

                return jsonify({
                    'success': True,
                    'records': records or [],
                    'current_page': page,
                    'total_pages': total_pages,
                    'total_items': total,
                    'per_page': per_page  # 返回当前每页显示数量
                })
        except Exception as e:
            print(f"Database error: {str(e)}")
            return jsonify({
                'success': False,
                'message': str(e),
                'records': [],
                'current_page': 1,
                'total_pages': 1,
                'total_items': 0
            })
        finally:
            connection.close()
    except Exception as e:
        print(f"Error in get_data: {str(e)}")
        return jsonify({
            'success': False,
            'message': f"Internal server error: {str(e)}",
            'records': [],
            'current_page': 1,
            'total_pages': 1,
            'total_items': 0
        }), 500

# ... (保持其他路由不变)

@train_bp.route('/api/train/delete_data/<int:id>', methods=['DELETE'])
def delete_data(id):
    try:
        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                # 检查记录是否存在
                cursor.execute("SELECT id FROM new_abalone WHERE id = %s", (id,))
                if not cursor.fetchone():
                    return jsonify({'success': False, 'message': '记录不存在'}), 404

                # 删除记录
                cursor.execute("DELETE FROM new_abalone WHERE id = %s", (id,))
            connection.commit()

            return jsonify({'success': True, 'message': '删除成功'})
        finally:
            connection.close()
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500


@train_bp.route('/api/train/retrain', methods=['POST'])
def retrain():
    global current_model, model_type, X_train, X_test, y_train, y_test

    try:
        data = request.json
        model_type = data.get('model_type', 'LR')

        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                # 从数据库加载数据
                cursor.execute("""
                SELECT sex, length, diameter, height, whole_weight, 
                       shucked_weight, viscera_weight, shell_weight, rings
                FROM new_abalone
                """)
                records = cursor.fetchall()

            if not records:
                return jsonify({'success': False, 'message': '没有数据可用于训练'}), 400

            # 准备数据
            data_list = []
            for record in records:
                data_list.append([
                    1 if record['sex'] == 'M' else (0 if record['sex'] == 'F' else -1),
                    record['length'],
                    record['diameter'],
                    record['height'],
                    record['whole_weight'],
                    record['shucked_weight'],
                    record['viscera_weight'],
                    record['shell_weight'],
                    record['rings']
                ])

            df = pd.DataFrame(data_list, columns=[
                'sex', 'length', 'diameter', 'height', 'whole_weight',
                'shucked_weight', 'viscera_weight', 'shell_weight', 'rings'
            ])

            # 特征和目标
            X = df.drop('rings', axis=1)
            y = df['rings']

            # 划分训练集和测试集
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.3, random_state=42)

            # 训练模型
            start_time = time.time()

            if model_type == 'LR':
                model = LinearRegression()
            elif model_type == 'DT-R':
                model = DecisionTreeRegressor(random_state=42)
            elif model_type == 'RF-R':
                model = RandomForestRegressor(n_estimators=100, random_state=42)
            else:
                return jsonify({'success': False, 'message': '不支持的模型类型'}), 400

            model.fit(X_train, y_train)
            current_model = model

            # 评估模型
            y_pred = model.predict(X_test)
            mae = mean_absolute_error(y_test, y_pred)
            mse = mean_squared_error(y_test, y_pred)
            rmse = np.sqrt(mse)
            r2 = r2_score(y_test, y_pred)

            train_time = time.time() - start_time

            return jsonify({
                'success': True,
                'model_type': model_type,
                'mae': mae,
                'mse': mse,
                'rmse': rmse,
                'r2': r2,
                'train_size': len(X_train),
                'test_size': len(X_test),
                'train_time': round(train_time, 6)
            })
        finally:
            connection.close()
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500